{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d55a4d45",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'torch'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Input \u001b[0;32mIn [2]\u001b[0m, in \u001b[0;36m<cell line: 4>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdata_process\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdata_iter\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DataIter\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m# from model.lstm import Seq2Seq\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnn\u001b[39;00m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'torch'"
     ]
    }
   ],
   "source": [
    "from data_process.sort_process import SortProcess\n",
    "from data_process.data_iter import DataIter\n",
    "# from model.lstm import Seq2Seq\n",
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7f483629",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "30"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(sp.c2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9493366",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "d4332884",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.7907199020738953\n",
      "1.9718694421980116\n",
      "1.394076371634448\n",
      "0.8911760791584298\n",
      "0.5398531026310391\n",
      "0.3316260027112784\n",
      "0.21492015653186375\n",
      "0.14373017419819478\n",
      "0.09693553271117034\n",
      "0.07084569838587884\n"
     ]
    }
   ],
   "source": [
    "src_path = \"data/sort/letters_source.txt\"\n",
    "trg_path = \"data/sort/letters_target.txt\"\n",
    "\n",
    "sp = SortProcess(src_path, trg_path)\n",
    "src_ids, trg_ids, label_ids = sp.dataset(max_len=10, to_numpy=False)\n",
    "train_iter = DataIter(list(zip(src_ids, trg_ids, label_ids)), batch_size=32, split=True)\n",
    "\n",
    "model = Seq2Seq(encode_vocab_size=len(sp.c2i),\n",
    "                decoder_vocab_size=len(sp.c2i),\n",
    "                emd_dim=64,\n",
    "                hid_size=128,\n",
    "                n_layers=2)\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
    "\n",
    "# for i in range(10):\n",
    "#     loss_list = []\n",
    "#     for i, (src, tti, tto) in enumerate(train_iter):\n",
    "#         src = torch.LongTensor(src)\n",
    "#         tti = torch.LongTensor(tti)\n",
    "#         tto = torch.LongTensor(tto)\n",
    "#         optimizer.zero_grad()\n",
    "#         output = model.forward(src, tti)\n",
    "\n",
    "#         tto = tto.reshape(-1)\n",
    "#         output = output.reshape(-1, output.shape[-1])\n",
    "#         loss = criterion(output, tto)\n",
    "#         loss.backward()\n",
    "#         optimizer.step()\n",
    "#         loss_list.append(float(loss.detach().numpy()))\n",
    "#     print(f\"{np.mean(loss_list)}\")\n",
    "    \n",
    "\n",
    "for i in range(10): \n",
    "    train_iter2 = tain_iter()\n",
    "    loss_list = []\n",
    "    for i, (src, tti, tto) in enumerate(train_iter2):\n",
    "        src = torch.LongTensor(src)\n",
    "        tti = torch.LongTensor(tti)\n",
    "        tto = torch.LongTensor(tto)\n",
    "        optimizer.zero_grad()\n",
    "        output = model.forward(src, tti)\n",
    "\n",
    "        tto = tto.reshape(-1)\n",
    "        output = output.reshape(-1, output.shape[-1])\n",
    "        loss = criterion(output, tto)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        loss_list.append(float(loss.detach().numpy()))\n",
    "    print(f\"{np.mean(loss_list)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cb1fc53",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "1.6240107748455133\n",
    "0.2903840807489694\n",
    "0.07381353720118063\n",
    "0.03429203448644556\n",
    "0.015125151696844032\n",
    "0.008991779296840438\n",
    "0.006034049688996122\n",
    "0.0042958749791446585\n",
    "0.0031844372451959994\n",
    "0.0024223972480898848\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8fcc49e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e25245b6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21bee4b9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f01a2464",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8df3c71",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "d53c7545",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.3394079208374023\n",
      "1.8272887468338013\n",
      "1.1600987911224365\n",
      "0.7091215252876282\n",
      "0.4507371187210083\n",
      "0.2776341140270233\n",
      "0.1770036816596985\n",
      "0.11481419205665588\n",
      "0.08303409814834595\n",
      "0.06585372239351273\n"
     ]
    }
   ],
   "source": [
    "model = Seq2Seq(encode_vocab_size=len(data_c2i),\n",
    "                decoder_vocab_size=len(taget_c2i),\n",
    "                emd_dim=64,\n",
    "                hid_size=128,\n",
    "                n_layers=2)\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
    "\n",
    "for i in range(10):\n",
    "    train_iter = tain_iter()\n",
    "    for i, (src, tti, tto) in enumerate(train_iter):\n",
    "        src = torch.LongTensor(src)\n",
    "        tti = torch.LongTensor(tti)\n",
    "        tto = torch.LongTensor(tto)\n",
    "        optimizer.zero_grad()\n",
    "        output = model.forward(src, tti)\n",
    "\n",
    "        tto = tto.reshape(-1)\n",
    "        output = output.reshape(-1, output.shape[-1])\n",
    "        loss = criterion(output, tto)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(f\"{loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66b06fc0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ed13584e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "# src_path = \"data/sort/letters_source.txt\"\n",
    "# trg_path = \"data/sort/letters_target.txt\"\n",
    "\n",
    "data_path = \"data/sort/letters_source.txt\"\n",
    "target_path = \"data/sort/letters_target.txt\"\n",
    "with open(data_path, 'r') as f:\n",
    "    data_s = f.read().split('\\n')\n",
    "    \n",
    "with open(target_path, 'r') as f:\n",
    "    target_s = f.read().split('\\n')\n",
    "    \n",
    "data_chars = sorted(list(set(\"\".join(data_s))))\n",
    "taget_chars = sorted(list(set(\"\".join(target_s))))\n",
    "pad_flage = '<PAD>'\n",
    "unk_flage = '<UNK>'\n",
    "beg_flage = '<GO>'\n",
    "end_flage = '<EOS>'\n",
    "data_chars = [pad_flage, unk_flage, beg_flage, end_flage] + data_chars\n",
    "taget_chars = [pad_flage, unk_flage, beg_flage, end_flage] + taget_chars\n",
    "\n",
    "data_c2i = {char: i for i, char in enumerate(data_chars)}\n",
    "taget_c2i = {char: i for i, char in enumerate(taget_chars)}\n",
    "taget_i2c = {i: char for i, char in enumerate(taget_chars)}\n",
    "\n",
    "max_len = 10\n",
    "\n",
    "data = []\n",
    "for line_ in data_s:\n",
    "    line = [data_c2i.get(char, 1) for char in  line_]\n",
    "    line = line[:max_len]\n",
    "    line = line + [0]*(max_len - len(line))\n",
    "    data.append(line)\n",
    "target_i = []\n",
    "target_o = []\n",
    "for line_ in target_s:\n",
    "    line = [2] + [taget_c2i.get(char, 1) for char in  line_] + [3]\n",
    "    line = line[:max_len]\n",
    "    line = line + [0]*(max_len - len(line))\n",
    "    \n",
    "    target_i.append(line)\n",
    "    line = [taget_c2i.get(char, 1) for char in  line_] + [3]\n",
    "    line = line[:max_len]\n",
    "    line = line + [0]*(max_len - len(line))\n",
    "    target_o.append(line)\n",
    "    \n",
    "    \n",
    "data = np.array(data)\n",
    "target_i = np.array(target_i)\n",
    "target_o = np.array(target_o)\n",
    "\n",
    "split_index = int(len(data)*0.7)\n",
    "train_data = data[:split_index]\n",
    "train_target_i = target_i[:split_index]\n",
    "train_target_o = target_o[:split_index]\n",
    "\n",
    "test_data = data[split_index:]\n",
    "target_data = target_o[split_index:]\n",
    "\n",
    "def tain_iter(batch_size=128):\n",
    "    num = len(train_data)//batch_size\n",
    "    for i in range(num):\n",
    "        beg_index = i*batch_size\n",
    "        end_index = (i+1)*batch_size\n",
    "        td = train_data[beg_index: end_index]\n",
    "        tti = train_target_i[beg_index: end_index]\n",
    "        tto = train_target_o[beg_index: end_index]\n",
    "        yield td, tti, tto\n",
    "\n",
    "def test_iter(batch_size=128):\n",
    "    num = len(test_data)//batch_size\n",
    "    for i in range(num):\n",
    "        beg_index = i*batch_size\n",
    "        end_index = (i+1)*batch_size\n",
    "\n",
    "        yield test_data[beg_index: end_index], target_data[beg_index: end_index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5db6e7d2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
